import GAN
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
import matplotlib.pyplot as plt
import time


class MNIST_DCGAN(object):

    def __init__(self):
        self.img_rows = 28
        self.img_cols = 28
        self.channel = 1
        self.x_train = input_data.read_data_sets(
            "mnist", one_hot=True).train.images
        self.x_train = self.x_train.reshape(-1, self.img_rows,
                                            self.img_cols, 1).astype(np.float32)

        self.DCGAN = GAN.DCGAN(28, 28, 1)
        self.discriminator = self.DCGAN.discriminator_model()
        self.adversarial = self.DCGAN.adversarial_model()
        self.generator = self.DCGAN.generator()

    def train(self, train_steps=2000, batch_size=256):
        d = []
        a = []
        for i in range(train_steps):
            t = time.time()
            images_train = self.x_train[np.random.randint(
                0, self.x_train.shape[0], size=batch_size), :, :, :]
            noise = np.random.uniform(-1.0, 1.0, size=[batch_size, 100])
            images_fake = self.generator.predict(noise)
            x = np.concatenate((images_train, images_fake))
            y = np.ones([2 * batch_size, 1])
            y[batch_size:, :] = 0
            for j in range(5):
                d_loss = self.discriminator.train_on_batch(x, y)
            y = np.ones([batch_size, 1])
            noise = np.random.uniform(-1.0, 1.0, size=[batch_size, 100])
            a_loss = self.adversarial.train_on_batch(noise, y)
            if i % 100 == 0:
                self.plot(i)
            t = time.time() - t
            if train_steps <= 200 or i >= (train_steps-200):
                plt.figure(1)
                #plt.axis([0, train_steps, -0.1, 1])
                plt.ion()
            d.append(d_loss[0])
                plt.plot(range(i + 1), d, linestyle='-', color='c')
                plt.grid(True, linestyle='--', color='g',linewidth='0.1')
                plt.figure(2)
                #plt.axis([0, train_steps, -0.1, 1.5])
                plt.ion()
                a.append(a_loss[0])
                plt.plot(range(i + 1), a, linestyle='-', color='c')
                plt.grid(True, linestyle='--', color='g',linewidth='0.1')
                plt.pause(0.05)
        plt.figure(1)
        plt.savefig('d_loss.jpg',dpi=200)
        plt.figure(2)
        plt.savefig('a_loss.jpg',dpi=200)
        plt.close()

    def plot(self, index):
        noise = np.random.uniform(-1.0, 1.0, size=[16, 100])
        images = self.generator.predict(noise)
        plt.figure(figsize=(10, 10))
        for i in range(images.shape[0]):
            plt.subplot(4, 4, i + 1)
            image = images[i, :, :, :]
            image = np.reshape(image, [self.img_rows, self.img_cols])
            plt.imshow(image, cmap='gray')
            plt.axis('off')
        plt.tight_layout()
        filename = 'mnist_%d.png' % index
        plt.savefig(filename)
        plt.close()
